"""
Langgraph 处理任务流
"""
from typing import List

from langchain_core.messages import HumanMessage, AIMessage
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from pydantic import BaseModel, Field

from model import get_model
from tools import get_weather

from langgraph.graph import StateGraph  # 整个Graph的图
from langgraph.graph import END  # 结束的节点
from langgraph.prebuilt import ToolNode  # 工具节点
from langgraph.graph import MessagesState  # 节点之间设置传递的状态值
from langgraph.checkpoint.memory import MemorySaver  # 保存节点之间所有消息,类似于history (回退)
from langgraph.prebuilt import chat_agent_executor

# --------------- 绑定tools --------------
tools = [get_weather]
model = get_model()
model = model.bind_tools(tools)


# 以下的功能提示词，可以参考智谱文档描述：https://www.bigmodel.cn/dev/howuse/prompt
# --------------- 制定多计划的节点 --------------
class Plans(BaseModel):
    plans: List[str] = Field("如果没有任务可划分,该内容为空")


take_plans_prompt = ChatPromptTemplate.from_messages(
    [
        (
            # 定义系统消息，为生成计划提供规则和要求
            "system",
            """对于给定的目标，提出一个简单的逐步计划。这个计划应该包含独立的任务，
            如果正确执行将得出正确的答案。不要添加任何多余的步骤。最后一步的结果应
            该是最终答案。确保每一步都有所有必要的信息 - 不要跳过步骤。""",
        ),
        # 定义占位符消息，用于后续填充实际的用户输入
        ("user", "{messages}"),
    ]
)


def take_many_plans(state: MessagesState):
    messages = state['messages']
    plans_chain = take_plans_prompt | model.with_structured_output(schema=Plans)
    response = plans_chain.invoke({'messages': messages})
    print("=" * 20, "Agent1", "=" * 20)
    print(response)
    tasks = response.__dict__['plans']
    # 例如该结果为：plans=['确定中国的直辖市', '查找每个直辖市的今天天气', '整理并输出结果']
    return {"messages": [HumanMessage(tasks)]}


#  --------------- 重新制定任务节点 --------------

replanner_prompt = ChatPromptTemplate.from_template(
    """
        对于给定的目标，提出一个简单的逐步计划，步骤必须越少越好。这个计划应该包含独立的任务，如果正确执行将得出正确的答案。不要添加任何多余的步骤。最后一步的结果应该是最终答案。确保每一步都有所有必要的信息 - 不要跳过步骤。
        你的目标是：
        {input}
        你的原计划是：
        {tasks}
        你目前已完成的步骤是：
        {finish}
        返回剩下的步骤，不要返回无原计划不相关内容，如果没有剩余步骤，则返回“END”
    """
)
# 相应地更新你的计划。如果不需要更多步骤并且可以返回给用户，那么就这样回应。如果需要，填写计划。只添加仍然需要完成的步骤。不要返回已完成的步骤作为计划的一部分。


def re_take_many_plans(state: MessagesState):
    print("=" * 20, "Agents3收到的信息", "=" * 20)
    input = state['messages'][0].content
    tasks = state['messages'][1].content
    tasks = "\n".join(tasks)
    finish = state['messages'][-2].content
    print("==当前目标",input)
    print("==任务信息",tasks)
    print("==已完成的信息",finish)
    plans_chain = replanner_prompt | model.with_structured_output(schema=Plans)
    response = plans_chain.invoke({"input": input, "tasks": tasks, "finish": finish})
    print("=" * 20, "Agent3", "=" * 20)
    tasks = response.__dict__['plans']
    print(tasks)
    return {"messages": [HumanMessage(tasks)]}


# --------------- 制定单任务执行的节点 --------------


def take_one_task(state: MessagesState):
    # print("=" * 20, "Agent2收到的信息", "=" * 20)
    # print(state['messages'])
    tasks = state['messages'][-1].content
    first_task = tasks[0]
    response = model.invoke([HumanMessage(first_task)])
    return {"messages": [response]}


def should_continue(state: MessagesState):  # 设置分支条件边
    messages = state['messages']
    last_message = messages[-1]
    if last_message.tool_calls:  # should_continue 直接返回节点
        return "tools"
    return "agent3"


def should_continue2(state: MessagesState):  # 设置分支条件边
    messages = state['messages']
    last_message = messages[-1]
    if "END" in last_message.content:  # should_continue 直接返回节点
        return END
    return "agent2"


tools_node = ToolNode(tools)
workflow = StateGraph(MessagesState)
workflow.add_node("agent1", take_many_plans)
workflow.add_node("agent2", take_one_task)
workflow.add_node("agent3", re_take_many_plans)
workflow.add_node("tools", tools_node)
workflow.set_entry_point("agent1")  # 设置起点位置

workflow.add_edge("agent1", "agent2")
workflow.add_edge("tools", "agent3")

workflow.add_conditional_edges(
    "agent2",
    should_continue,
    {
        "tools": "tools",
        "agent3": "agent3"
    }
)

workflow.add_conditional_edges(
    "agent3",
    should_continue2,
)
app = workflow.compile(checkpointer=MemorySaver())  # 构建Langgraph应用

final_state = app.invoke(
    {"messages": [HumanMessage("重庆的天气，上海的天气")]},
    config={"configurable": {"thread_id": 42}}
)
result = final_state["messages"][-1].content
print(result)
